/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */

package com.pictavore.image.fits;

import java.awt.image.*;
import java.io.*;
import java.awt.*;
import java.util.*;
import java.lang.Object.*;
import java.nio.*;
import javax.media.jai.*;

import com.pictavore.util.Utils;

/**
 *
 * @author dvenable
 */
public class Fits {
    
    Hashtable<String, String> dict = new Hashtable<String, String>();
    Hashtable<String, String> cmts = new Hashtable<String, String>();
    int nbits = 0;
    int npix = 0;
    int nlin = 0;
    int len = 0;
    double bzero = 0;
    double bscale = 1.0;
    byte[] raw = null;
    int minval = Integer.MAX_VALUE;
    int maxval = Integer.MIN_VALUE;
    int rawmin = Integer.MAX_VALUE;
    int rawmax = Integer.MIN_VALUE;
    int[] udata = null;
    String info = null;
    
    public int getNBits() { return nbits; }
    public int getWidth() { return npix; }
    public int getHeight() { return nlin; }
    public int getNumPix() { return len; }
    public double getBZero() { return bzero; }
    public double getBScale() { return bscale; }
    public byte[] getRawData() { return raw; }
    public int getMinValue() { return minval; }
    public int getMaxValue() { return maxval; }
    public int getRawMin() { return rawmin; }
    public int getRawMax() { return rawmax; }
    public int[] getUnsignedData() { return udata; }
    public String getInfo() { return info; }


    public static Boolean IsFits(String path)
    {
        return (Utils.getExtension(path).equalsIgnoreCase("fits"));
    }


    public Fits(String srcPath) throws FileNotFoundException, IOException, FITSException
    {
        BufferedInputStream is = new BufferedInputStream(new FileInputStream(srcPath));


        // <editor-fold defaultstate="collapsed" desc="Read the FITS header">
        boolean loopAgain = true;
        byte[] hdrBytes = new byte[2880];
        while (loopAgain)
        {
            // read the first header chunk
            int nb = is.read(hdrBytes);
            if (nb == -1)
                throw new FITSException("Invalid FITS header");
            
            String hdr = new String(hdrBytes);

            // scan thru each card
            for (int k = 0, kk = 0; k < 36; ++k, kk += 80)
            {
                // get the card
                String card = hdr.substring(kk, kk + 80);

                // get the key
                String key = card.substring(0, 8).trim();

                // is a value specified?
                boolean isval = (card.charAt(8) == '=') && (card.charAt(9) == ' ');

                String value = null;
                String comment = null;
                if (isval)
                {
                    // get the value (num or string)
                    if (card.charAt(10) == '\'')
                    {
                        int ind = card.indexOf('\'', 11);
                        if (ind == -1)
                            throw new FITSException("Illegal string in FITS header");
                        value = card.substring(11, ind);
                        comment = card.substring(ind + 1).trim();
                    }
                    else
                    {
                        value = card.substring(10, 30).trim();
                        comment = card.substring(30).trim();
                    }
                }
                else
                    comment = card.substring(8,80).trim();
                
                // store header entried into dictionaries
                if (comment != null && comment.length() > 0)
                    cmts.put(key, comment);
                
                if (value != null)
                    dict.put(key, value);

                // quit scanning if we find an end card
                if (key.equals("END"))
                { loopAgain = false; break; }

                // kinda ignore this :)
//                if (key.equals("EXTEND"))
//                    loopAgain = true;

            }
        }
        // </editor-fold>

        //<editor-fold desc="parse relevant header info">
        // a few diagnostics
        if (dict.containsKey("SIMPLE")==false || dict.get("SIMPLE").equals("T") == false)
            throw new FITSException("Not a SIMPLE FITS file");
        
        if (dict.containsKey("NAXIS")==false || dict.get("NAXIS").equals("2") == false)
            throw new FITSException("Can only read 2D images");
        
        if (dict.containsKey("BITPIX")==false)
            throw new FITSException("#bits not specified");
        else
            nbits = Integer.parseInt(dict.get("BITPIX"));
        
        if (dict.containsKey("NAXIS1")==false)
            throw new FITSException("Image width not specified");
        else
            npix = Integer.parseInt(dict.get("NAXIS1"));
        
        if (dict.containsKey("NAXIS2")==false)
            throw new FITSException("Image height not specified");
        else
            nlin = Integer.parseInt(dict.get("NAXIS2"));
        
        if (dict.containsKey("BZERO")==false)
            bzero = 0.0;
        else
            bzero = Double.parseDouble(dict.get("BZERO"));
        
        if (dict.containsKey("BSCALE")==false)
            bscale = 1.0;
        else
            bscale = Double.parseDouble(dict.get("BSCALE"));
        
        if (nbits == -16 || nbits == -32)
            throw new FITSException("Floating point formats not supported");

        info = "";
        for(String kk : dict.keySet())
        {
            String tmp = String.format("  %s: %s  /%s\n", kk, dict.get(kk), cmts.containsKey(kk)?cmts.get(kk):"");
            info = info.concat(tmp);
        }
        System.out.println(info);
        //</editor-fold>
        
        // <editor-fold defaultstate="collapsed" desc="Fetch the raw image data">
        len = npix * nlin;
        int blen = len * (nbits/8);
        raw = new byte[blen];
        int nb = is.read(raw);
        if (nb != blen)
            System.out.println("Did not read enough data bytes");
        
        // calc min and max values
        udata = new int[len];
        if (nbits == 8)
        {
            for (int k=0; k<len; ++k)
            {
                int p = (raw[k] & 0x00FF);
                rawmin = (p<rawmin) ? p : rawmin;
                rawmax = (p>rawmax) ? p : rawmax;

                p = (int)(bzero + bscale*p);
                udata[k] = p;
                minval = (p<minval) ? p : minval;
                maxval = (p>maxval) ? p : maxval;
            }
        }
        else if (nbits == 16)
        {
            ByteBuffer bb = ByteBuffer.wrap(raw);
            bb.order(ByteOrder.BIG_ENDIAN);
            for (int k=0; k<len; ++k)
            {
                short v = bb.getShort();
                rawmin = (v<rawmin) ? v : rawmin;
                rawmax = (v>rawmax) ? v : rawmax;

                int p = (int)(bzero + bscale*(double)v);
                udata[k] = p;
                maxval = (p > maxval) ? p : maxval;
                minval = (p < minval) ? p : minval;
           }
        }
        else if (nbits == 32)
        {
            ByteBuffer bb = ByteBuffer.wrap(raw);
            bb.order(ByteOrder.BIG_ENDIAN);
            for (int k=0; k<len; ++k)
            {
                int v = bb.getInt();
                rawmin = (v<rawmin) ? v : rawmin;
                rawmax = (v>rawmax) ? v : rawmax;

                int p = (int)(bzero + bscale*(double)v);
                udata[k] = p;
                maxval = (p > maxval) ? p : maxval;
                minval = (p < minval) ? p : minval;
            }
        }
        // release raw data
        //raw = null;

        // </editor-fold>  
        
        // done acquiring the image data
        is.close();
    }

    
    public BufferedImage getBufferedImage()
    {
        BufferedImage bimg = null;
        
        if (nbits == 8)
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_BYTE_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferByte usbuf = (DataBufferByte) rst.getDataBuffer();
            byte[] pix = (byte[]) usbuf.getData();
            for (int k=0; k<len; ++k)
                pix[k] = (byte)(udata[k] & 0x000000FF);
            bimg.setData(rst);

        }
        else if (nbits == 16)
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_USHORT_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferUShort usbuf = (DataBufferUShort) rst.getDataBuffer();
            short[] pix = (short[]) usbuf.getData();
            for (int k=0; k<len; ++k)
                pix[k] = (short)(udata[k] & 0x0000FFFF);
            bimg.setData(rst);
        }
        else if (nbits == 32)
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_USHORT_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferUShort usbuf = (DataBufferUShort) rst.getDataBuffer();
            short[] pix = (short[]) usbuf.getData();
            for (int k=0; k<len; ++k)
                pix[k] = (short)((udata[k]>>16) & 0x0000FFFF);
            bimg.setData(rst);
        }
        return bimg;
    }

    public TiledImage getTiledImage()
    {

        TiledImage timg = null;

        if (nbits == 8)
        {
            byte[] buf = new byte[len];
            for (int k=0; k<len; ++k)
                buf[k] = (byte)(udata[k] & 0x000000FF);
            DataBufferByte db = new DataBufferByte(buf, len);
            SampleModel smp = RasterFactory.createBandedSampleModel(DataBuffer.TYPE_BYTE, npix, nlin, 1);
            ColorModel cm = PlanarImage.createColorModel(smp);
            WritableRaster rst = RasterFactory.createWritableRaster(smp, db, new Point(0,0));
            timg = new TiledImage(0,0,npix,nlin,0,0,smp,cm);
            timg.setData(rst);
            return timg;
        }
        else if (nbits == 16)
        {
            short[] buf = new short[len];
            for (int k=0; k<len; ++k)
                buf[k] = (short)(udata[k] & 0x0000FFFF);
            DataBufferUShort db = new DataBufferUShort(buf, len);
            SampleModel smp = RasterFactory.createBandedSampleModel(DataBuffer.TYPE_USHORT, npix, nlin, 1);
            ColorModel cm = PlanarImage.createColorModel(smp);
            WritableRaster rst = RasterFactory.createWritableRaster(smp, db, new Point(0,0));
            timg = new TiledImage(0,0,npix,nlin,0,0,smp,cm);
            timg.setData(rst);
            return timg;
        }
        else if (nbits == 32)
        {
            short[] buf = new short[len];
            for (int k=0; k<len; ++k)
                buf[k] = (short)(udata[k] & 0x0000FFFF);
            DataBufferUShort db = new DataBufferUShort(buf, len);
            SampleModel smp = RasterFactory.createBandedSampleModel(DataBuffer.TYPE_USHORT, npix, nlin, 1);
            ColorModel cm = PlanarImage.createColorModel(smp);
            WritableRaster rst = RasterFactory.createWritableRaster(smp, db, new Point(0,0));
            timg = new TiledImage(0,0,npix,nlin,0,0,smp,cm);
            timg.setData(rst);
            return timg;
      }
        return timg;
     }

    public BufferedImage getStretchedBufferedImage()
    {
        BufferedImage bimg = null;

        if (nbits == 8)
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_BYTE_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferByte usbuf = (DataBufferByte) rst.getDataBuffer();
            byte[] pix = (byte[]) usbuf.getData();
            double sf = (double)255 / (double)(maxval - minval);
            for (int k=0; k<len; ++k)
            {
                int p = (int)((double)(udata[k] - minval) * sf);
                pix[k] = (byte)(p & 0x000000FF);
            }
            bimg.setData(rst);
       }
        else if (nbits == 16)
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_USHORT_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferUShort usbuf = (DataBufferUShort) rst.getDataBuffer();
            short[] pix = (short[]) usbuf.getData();
            double sf = (double)65535 / (double)(maxval - minval);
            for (int k=0; k<len; ++k)
            {
                int p = (int)((double)(udata[k] - minval) * sf);
                pix[k] = (short)(p & 0x0000FFFF);
            }
            bimg.setData(rst);
        }
        else if (nbits == 32)   // really, this  is the same as 16 bits because BufferedImage has no 32bit iamge model.
        {
            bimg = new BufferedImage(npix, nlin, BufferedImage.TYPE_USHORT_GRAY);
            WritableRaster rst = bimg.getRaster();
            DataBufferUShort usbuf = (DataBufferUShort) rst.getDataBuffer();
            short[] pix = (short[]) usbuf.getData();
            double sf = (double)65535 / (double)(maxval - minval);
            for (int k=0; k<len; ++k)
            {
                int p = (int)((double)(udata[k] - minval) * sf);
                pix[k] = (byte)(p & 0x0000FFFF);
            }
            bimg.setData(rst);
        }
        return bimg;
    }

    public BufferedImage getStretchedImage(Rectangle win) throws FITSException
    {
        BufferedImage res = null;
        if (nbits == 8)
        {
            throw new FITSException("Stretching an 8bit window not implemented");
        }
        else if (nbits == 16)
        {
            res = new BufferedImage(win.width, win.height, BufferedImage.TYPE_USHORT_GRAY);
            WritableRaster rst = res.getRaster();
            DataBufferUShort usbuf = (DataBufferUShort) rst.getDataBuffer();
            short[] pix = (short[]) usbuf.getData();

            // calc min/max
            int W = getWidth();
            int H = getHeight();
            int w = win.width;
            int h = win.height;
            int min = 65535;
            int max = 0;
            int ofs = win.y*W + win.x;
            for (int y=0; y<h; ++y, ofs += W)
                for (int x=0, k=ofs; x<w; ++x, ++k)
                {
                    int p = udata[k];
                    min = (p<min) ? p : min;
                    max = (p>max) ? p : max;
                }

            // stretch to min/max
            double sf = (double)65535 / (double)(max - min);
            ofs = win.y*W + win.x;
            for (int y=0, kk=0; y<h; ++y, ofs += W)
                for (int x=0, k=ofs; x<w; ++x, ++k, ++kk)
                {
                    int p = (int)((double)(udata[k] - min) * sf);
                    pix[kk] = (short)(p & 0x0000FFFF);
                }
            res.setData(rst);
        }
        else if (nbits == 32)
        {
            throw new FITSException("Stretching an 32bit window not implemented");
        }
        return res;
    }

}

